import argparse
import copy
import json
import os
import sys
sys.path.append(".")
try:
    import apex
except:
    print("No APEX!")
import numpy as np
import torch
import yaml
from det3d import torchie
from det3d.datasets import build_dataloader, build_dataset
from det3d.models import build_detector
from det3d.torchie import Config
from det3d.torchie.apis import (
    batch_processor,
    build_optimizer,
    get_root_logger,
    init_dist,
    set_random_seed,
    train_detector,
)
from det3d.torchie.trainer import load_checkpoint
import pickle 
import time 
from matplotlib import pyplot as plt 
from det3d.torchie.parallel import collate, collate_kitti
from torch.utils.data import DataLoader
import matplotlib.cm as cm
import subprocess
import cv2
from tools.demo_utils import visual 
from collections import defaultdict

import json
from point_to_bev import *
from util import *

class_names=['car','truck','forklift','tractor','flatbed','trailer','bicycle','pedestrian','traffic_cone','water_horse','cyclist','guardrail','unknown']

def convert_box(info):
    boxes =  info["gt_boxes"].astype(np.float32)
    names = info["gt_names"]

    assert len(boxes) == len(names)

    detection = {}

    detection['box3d_lidar'] = boxes.tolist()

    # dummy value 
    detection['label_preds'] = [class_names.index(name) for name in names]
    detection['scores'] = np.ones(len(boxes)).tolist()

    return detection 

def main():
    cfg = Config.fromfile('custom_config/rcs16.py')
    
    model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)

    dataset = build_dataset(cfg.data.val)

    data_loader = DataLoader(
        dataset,
        batch_size=1,
        sampler=None,
        shuffle=False,
        num_workers=8,
        collate_fn=collate_kitti,
        pin_memory=False,
    )

    checkpoint = load_checkpoint(model, 'work_dirs/work_dirs_0625/rcs16/latest.pth', map_location="cpu")
    model.eval()

    model = model.cuda()

    cpu_device = torch.device("cpu")

    points_list = [] 
    gt_annos = [] 
    detections  = [] 

    for i, data_batch in enumerate(data_loader):
        scene_name=data_batch['metadata'][0]['token'].split("-")[1]
        # if '0408' not in scene_name:
        #     continue
        print("data index:",i)
        # if i %20!=0:
        #     continue
        # if i>10000:
        #     break
        info = dataset._nusc_infos[i]
        gt_annos.append(convert_box(info))

        points_ = data_batch['points'][0][:, 0:4].cpu().numpy()
        # points[:,[1,2]]=points[:,[2,1]]
        with torch.no_grad():
            outputs = batch_processor(
                model, data_batch, train_mode=False, local_rank=0,
            )
        for output in outputs:
            for k, v in output.items():
                if k not in [
                    "metadata",
                ]:
                    output[k] = v.to(cpu_device).numpy().tolist()
                else:
                    output[k]['image_prefix'] = str(output[k]['image_prefix'])
            output["name"]=data_batch['metadata'][0]['token']
            detections.append(output)

        
        bev_image=point_to_bev(points_)
        # print("gt:")
        for obj in convert_box(info)['box3d_lidar']:
            x,y,z,l,w,h,vx,vy,yaw=obj
            # print(x,y,z,l,w,h,vx,vy,yaw)
            roll,pitch=0,0

            points=calculate_vertices(x, y, z,l,w, h,yaw)
            points[:,0]=1080/2+points[:,0]*1080/60
            points[:,1]=1920/2-points[:,1]*1920/60
            points=points.astype(np.int32)

            cv2.line(bev_image,(points[0][0],points[0][1]) , (points[1][0],points[1][1]) ,(0,0,255),1)
            cv2.line(bev_image,(points[1][0],points[1][1]) , (points[2][0],points[2][1]) ,(0,255,255),1)
            cv2.line(bev_image,(points[2][0],points[2][1]) , (points[3][0],points[3][1]) ,(0,0,255),1)
            cv2.line(bev_image,(points[3][0],points[3][1]) , (points[0][0],points[0][1]) ,(0,0,255),1)

        # print("pre:")
        for j in range(np.array(outputs[0]['box3d_lidar']).shape[0]):
            obj=outputs[0]['box3d_lidar'][j]
            score=outputs[0]['scores'][j]
            if score<0.5:
                continue
            x,y,z,l,w,h,vx,vy,yaw=obj
            # print(x,y,z,l,w,h,vx,vy,yaw)
            roll,pitch=0,0

            points=calculate_vertices(x, y, z,l,w, h,yaw)
            points[:,0]=1080/2+points[:,0]*1080/60
            points[:,1]=1920/2-points[:,1]*1920/60
            points=points.astype(np.int32)

            cv2.line(bev_image,(points[0][0],points[0][1]) , (points[1][0],points[1][1]) ,(0,255,0),1)
            cv2.line(bev_image,(points[1][0],points[1][1]) , (points[2][0],points[2][1]) ,(0,255,255),1)
            cv2.line(bev_image,(points[2][0],points[2][1]) , (points[3][0],points[3][1]) ,(0,255,0),1)
            cv2.line(bev_image,(points[3][0],points[3][1]) , (points[0][0],points[0][1]) ,(0,255,0),1)
        cv2.imwrite(f"demo_rcs/{data_batch['metadata'][0]['token']}.jpg",bev_image)

        points_list.append(points_.T)
        # points_list.append(points)

    json.dump(detections, open('pre_0624.json', 'w'))
    json.dump(gt_annos, open('gt_0624.json', 'w'))
    
    print('Done model inference. Please wait a minute, the matplotlib is a little slow...')
    
    # for i in range(len(points_list)):
    #     visual(points_list[i], gt_annos[i], detections[i], i)
    #     print("Rendered Image {}".format(i))
    
    # image_folder = 'demo'
    # video_name = 'video.avi'

    # images = [img for img in os.listdir(image_folder) if img.endswith(".png")]
    # images.sort(key=lambda img_name: int(img_name.split('.')[0][4:]))
    # frame = cv2.imread(os.path.join(image_folder, images[0]))
    # height, width, layers = frame.shape

    # video = cv2.VideoWriter(video_name, 0, 1, (width,height))
    # cv2_images = [] 

    # for image in images:
    #     cv2_images.append(cv2.imread(os.path.join(image_folder, image)))

    # for img in cv2_images:
    #     video.write(img)

    # cv2.destroyAllWindows()
    # video.release()

    print("Successfully save video in the main folder")

if __name__ == "__main__":
    main()
